from epics import PV
import time
import os
from pcaspy import SimpleServer, Driver
import threading
import numpy as np
import pandas as pd

mod_list = [["H1", "H2"], ["V1", "V2"]]
Power_DATAself = "iLINAC:mod1_FC1_Rd_Avg"  # 采集的数据PV名设定
Pos_read = "iLINAC:dActPos_Slit{}"
Start_Pos = "iLINAC:dStartPos_Slit{}"
Step = "iLINAC:dSetStep_{}"
End_Pos = "iLINAC:dEndPos_Slit{}"
Pos_Opr = "iLINAC:bCmdEmittanceMove_{}"
Acquisition_frequency = 10  # 数据采集频率：一秒钟采集多少数据，最多为100
EPICS_CA_ADDR_LIST = "127.0.0.1:5066 192.168.1.26 192.168.1.27"  # 采集数据PV所在服务器的IP地址，默认为localhost
EPICS_CAS_SERVER_PORT = "5066"
threshold = 1 / Acquisition_frequency


def pad_data_linearly(data, target_length):
    """
    将数据线性补齐到目标长度。

    参数:
    data (list or np.array): 原始数据，长度小于或等于target_length。
    target_length (int): 目标长度。

    返回:
    np.array: 补齐后的数据，长度为target_length。
    """

    # 计算新数据点的间隔
    x_original = np.linspace(0, 1, len(data["X"]))  # 原始数据的索引位置（归一化到0-1）
    x_new = np.linspace(0, 1, target_length)  # 新数据的索引位置（归一化到0-1）

    # 计算新数据点的间隔
    y_original = np.linspace(0, 1, len(data["Y"]))  # 原始数据的索引位置（归一化到0-1）
    y_new = np.linspace(0, 1, target_length)  # 新数据的索引位置（归一化到0-1）

    # 计算新数据点的间隔
    z_original = np.linspace(0, 1, len(data["Z"]))  # 原始数据的索引位置（归一化到0-1）
    z_new = np.linspace(0, 1, target_length)  # 新数据的索引位置（归一化到0-1）

    # 使用线性插值来生成新数据
    x_data_new = np.interp(x_new, x_original, data["X"])
    y_data_new = np.interp(y_new, y_original, data["Y"])
    z_data_new = np.interp(z_new, z_original, data["Z"])

    return x_data_new, y_data_new, z_data_new


class emittance(threading.Thread):
    def __init__(self, setParam) -> None:
        print("Emittance is run")
        self.Stop_Flag = True
        self.data = PV(Power_DATAself)
        threading.Thread.__init__(self)
        self.setParam = setParam
        self.index = -1
        self.pv_dict = {0: {}, 1: {}}
        self.pv_dict[0]["Pos_read"] = [PV(Pos_read.format(i)) for i in mod_list[0]]
        self.pv_dict[0]["Start_Pos"] = [PV(Start_Pos.format(i)) for i in mod_list[0]]
        self.pv_dict[0]["End_Pos"] = [PV(End_Pos.format(i)) for i in mod_list[0]]
        self.pv_dict[0]["Step"] = PV(Step.format(mod_list[0][0]))
        self.pv_dict[0]["x_old_min_pos"] = abs(
            self.pv_dict[0]["Start_Pos"][0].get() - self.pv_dict[0]["End_Pos"][0].get()
        )
        self.pv_dict[0]["y_old_min_pos"] = abs(
            self.pv_dict[0]["Start_Pos"][1].get() - self.pv_dict[0]["End_Pos"][1].get()
        )

        self.pv_dict[1]["Pos_read"] = [PV(Pos_read.format(i)) for i in mod_list[1]]
        self.pv_dict[1]["Start_Pos"] = [PV(Start_Pos.format(i)) for i in mod_list[1]]
        self.pv_dict[1]["End_Pos"] = [PV(End_Pos.format(i)) for i in mod_list[1]]
        self.pv_dict[1]["Step"] = PV(Step.format(mod_list[1][0]))
        self.pv_dict[1]["x_old_min_pos"] = abs(
            self.pv_dict[1]["Start_Pos"][0].get() - self.pv_dict[1]["End_Pos"][0].get()
        )
        self.pv_dict[1]["y_old_min_pos"] = abs(
            self.pv_dict[1]["Start_Pos"][1].get() - self.pv_dict[1]["End_Pos"][1].get()
        )

        self.pv_dict[0]["Pos_Opr"] = PV(
            Pos_Opr.format(mod_list[0][0]), callback=self.onValueChange
        )
        self.pv_dict[1]["Pos_Opr"] = PV(
            Pos_Opr.format(mod_list[1][0]), callback=self.onValueChange
        )

        try:
            self.start()
        except:
            print("Error: No data")

    def onValueChange(self, pvname=None, value=None, host=None, **kws):
        for index in range(len(mod_list)):
            if pvname == self.pv_dict[index]["Pos_Opr"].pvname:
                self.Flag = True
                self.index = index
                print("PV:{}\t已触发".format(pvname))

    def X_get(self):
        x_min = self.pv_dict[self.index]["Start_Pos"][0].get()
        x_max = self.pv_dict[self.index]["End_Pos"][0].get()
        x_temp = self.pv_dict[self.index]["Step"].get()
        result = []
        data = x_min
        while True:
            if data > x_max:
                break
            result.append(data)
            data += x_temp
        if x_max % x_temp != 0:
            result.append(x_max)
        return result

    def x_y_z_get(self, datas):
        x_num = self.X_get()
        result = []
        for i in x_num:
            _data = datas[abs(datas["X"] - i) <= 0.01]
            result.append(_data)
        return result

    def Y_num_get(self, datas):
        y = []
        for i in self.X_get():
            y.append(len(datas[abs(datas["X"] - i) <= 0.01]))
        return min(y), max(y)

    def X_num_get(self):
        return len(self.X_get())

    def reset_min_pos(self):

        pos_data = self.pv_dict[self.index]["Pos_read"][1].get()
        start_pos = self.pv_dict[self.index]["Start_Pos"][1].get()
        end_pos = self.pv_dict[self.index]["End_Pos"][1].get()
        if abs(pos_data - start_pos) < threshold:
            print("Y到达起始位置,重置end的距离")
            self.pv_dict[self.index]["y_old_min_pos"] = abs(start_pos - end_pos)

    def pos_calc(self):
        # 使用更具描述性的变量名
        pos_data = self.pv_dict[self.index]["Pos_read"][1].get()

        try:
            # 获取边界位置
            start_pos = self.pv_dict[self.index]["Start_Pos"][1].get()
            end_pos = self.pv_dict[self.index]["End_Pos"][1].get()
            current_pos = abs(pos_data - end_pos)
            y_old_min_pos = self.pv_dict[self.index]["y_old_min_pos"]

            # 分别检查每个条件
            is_within_range = start_pos <= pos_data <= end_pos
            is_new_min_pos = current_pos <= y_old_min_pos
            print("Y的当前位置:", pos_data)
            print(
                "距离截止点:{},上个点距截止点位置：{}".format(
                    current_pos, y_old_min_pos
                )
            )
            if is_within_range and is_new_min_pos:
                self.pv_dict[self.index]["y_old_min_pos"] = current_pos
                return True, pos_data
            else:
                return False, pos_data
        except KeyError as e:
            # 异常处理
            print(f"KeyError: {e}")
            return False, None

    def step_calc(self):
        # 定义阈值为可配置参数
        try:
            pos_data = self.pv_dict[self.index]["Pos_read"][0].get()
            start_pos = self.pv_dict[self.index]["Start_Pos"][0].get()
            end_pos = self.pv_dict[self.index]["End_Pos"][0].get()
            step = self.pv_dict[self.index]["Step"].get()

            current_pos = abs(pos_data - end_pos)
            x_old_min_pos = self.pv_dict[self.index]["x_old_min_pos"]
            print("X的当前位置:", pos_data)
            print(
                "距离截止点:{},上个点距截止点位置：{}".format(
                    current_pos, x_old_min_pos
                )
            )
            if current_pos <= x_old_min_pos:
                self.pv_dict[self.index]["x_old_min_pos"] = current_pos
            else:
                if abs(pos_data - end_pos) < threshold:
                    return True, False, pos_data

            if pos_data >= start_pos and pos_data <= end_pos:
                mod = abs(pos_data - start_pos) % step
                if (
                    mod <= threshold
                    or abs(mod - step) <= threshold
                    or abs(pos_data - end_pos) < threshold
                ):
                    return True, True, pos_data
            return False, True, pos_data
        except Exception as e:
            # 处理异常情况
            print(f"An error occurred: {e}")
            return False, False, pos_data

    def run(self):
        time.sleep(1)
        self.Flag = False
        while True:
            time.sleep(1)
            if self.Flag and not self.index == -1:
                result_x = []
                result_y = []
                result_z = []
                x_pos_flag, x_Stop_flag, x_pos_data = self.step_calc()
                y_pos_flag, y_pos_data = self.pos_calc()
                self.setParam(
                    "iLINAC:bCmdEmittanceStatus_{}".format(mod_list[self.index][0]), 1
                )
                if x_pos_flag:

                    print("X方向电机已就位")
                    self.setParam(
                        "iLINAC:bCmdEmittanceStatus_{}".format(mod_list[self.index][0]),
                        2,
                    )
                    Stop_flag = False
                    while True:
                        if Stop_flag and x_pos_flag == False and y_pos_flag == False:
                            break
                        print("等待Y电机就位")
                        self.reset_min_pos()
                        while x_pos_flag and y_pos_flag:
                            print("采集数据")
                            result_x.append(x_pos_data)
                            result_y.append(y_pos_data)
                            result_z.append(self.data.get())
                            time.sleep(1 / Acquisition_frequency)
                            x_pos_flag, x_Stop_flag, x_pos_data = self.step_calc()
                            y_pos_flag, y_pos_data = self.pos_calc()
                            # print(x_pos_data,y_pos_data,self.data.get())
                        if x_Stop_flag == False:
                            Stop_flag = True
                        x_pos_flag, x_Stop_flag, x_pos_data = self.step_calc()
                        y_pos_flag, y_pos_data = self.pos_calc()
                        time.sleep(1 / Acquisition_frequency)
                    filename = time.strftime("%Y%m%d%H%M%S", time.localtime())
                    with open(
                        "./output/{}{}.csv".format(filename, mod_list[self.index][0]),
                        "w",
                    ) as f:
                        f.write("X,Y,FC1_Rd_Avg\n")
                        for i in range(len(result_x)):
                            f.write(
                                "{:.3f},{},{}\n".format(
                                    result_x[i], result_y[i], result_z[i]
                                )
                            )
                    df = pd.DataFrame(
                        {
                            "X": result_x,
                            "Y": result_y,
                            "Z": result_z,
                        }
                    )
                    df["Z"] = (
                        (df["Z"] - df["Z"].min())
                        / (df["Z"].max() - df["Z"].min())
                        * 255
                    )
                    self.setParam(
                        "iLINAC:bCmdEmittance_ArrayWidth", int(self.X_num_get())
                    )
                    self.setParam(
                        "iLINAC:bCmdEmittance_ArrayHeight", int(self.Y_num_get(df)[1])
                    )
                    datas = []
                    columns = ["X", "Y", "Z"]
                    for i in self.x_y_z_get(df):
                        a = pad_data_linearly(i, self.Y_num_get(df)[1])
                        _data = []
                        for x, y, z in zip(a[0], a[1], a[2]):
                            _data.append([float(x), float(y), float(z)])
                        _data = pd.DataFrame(_data, columns=columns)
                        _data = _data.sort_values(by=["Y"], ascending=[False])
                        datas.append(_data)
                    result = []
                    for j in range(self.Y_num_get(df)[1]):
                        for i in datas:
                            # print(
                            #     "X= {},Y= {},Z= {}".format(
                            #         float(i["X"][j]), float(i["Y"][j]), float(i["Z"][j])
                            #     )
                            # )
                            result.append(i["Z"][j])
                    self.setParam("iLINAC:bCmdEmittance_ArrayData", result)
                    self.pv_dict[self.index]["Pos_Opr"].put(0)
                    self.Flag = False
                    self.setParam(
                        "iLINAC:bCmdEmittanceStatus_{}".format(mod_list[self.index][0]),
                        0,
                    )
                    print("采集完成")


class myDriver(Driver):
    def __init__(self):
        super(myDriver, self).__init__()
        emittance(self.setParam)


if __name__ == "__main__":
    if not os.path.exists("./output"):
        os.mkdir("./output")
    os.environ["EPICS_CA_ADDR_LIST"] = EPICS_CA_ADDR_LIST
    os.environ["EPICS_CA_AUTO_ADDR_LIST"] = "no"
    os.environ["EPICS_CAS_SERVER_PORT"] = EPICS_CAS_SERVER_PORT
    os.environ["EPICS_CA_MAX_ARRAY_BYTES"] = "100000000"
    prefix = ""
    pvdb = {
        "iLINAC:bCmdEmittanceStatus_H1": {
            "type": "enum",
            "enums": ("Stop", "Wait", "Run"),
            "scan": 1,
        },
        "iLINAC:bCmdEmittanceStatus_V1": {
            "type": "enum",
            "enums": ("STOP", "Wait", "Run"),
            "scan": 1,
        },
        "iLINAC:bCmdEmittance_ArrayData": {
            "prec": 3,
            "count": 20000,
            "scan": 1,
        },
        "iLINAC:bCmdEmittance_ArrayWidth": {
            "prec": 0,
            "scan": 1,
        },
        "iLINAC:bCmdEmittance_ArrayHeight": {
            "prec": 0,
            "scan": 1,
        },
    }
    server = SimpleServer()
    server.createPV(prefix, pvdb)
    driver = myDriver()
    while True:
        server.process(0.1)
